#include "cktcover.h"
#include "rewrite/npn.h"
#include "gzstream.h"
#include "satsolver.h"

static uint8_t swap01_[4] = {0,2,1,3};
inline uint8_t swap01( uint8_t a ) { return swap01_[a]; }
Lit l0;
static bool slim( vec<Lit> & v )
{
   unsigned j=0;
   for( unsigned i = 0 ; i < v.size() ; ++ i )
   {
      if( v[i] == ~l0 )return false;
      else if( v[i] != l0  )v[j++] = v[i];
   }
   v.sz = j;
   return true;
}
struct BestCut
{
   unsigned short size :4 ;
   unsigned short idx : 12;
   float    cost;

   BestCut():size(2),idx(0),cost(0){}
   BestCut( unsigned s , unsigned i , float c ):size(s),idx(i),cost(c){}
};
/// LOCAL CLASS FOR WRITING CNF TO OGZ file
class SolverFile : public SATSolver
{
   ogzstream out;     
   public:
   SolverFile( const char * fname , unsigned nvar, unsigned ncls )
   {
      out.open( fname );
      out << "p cnf " <<nvar<< " " <<ncls<<"\n";
   }
   virtual ~SolverFile(){};
   virtual void addClause( vec<Lit> & c )
   {
      for( unsigned i = 0 ; i < c.size() ; ++i )
         out << toDimacs( c[i] ) << ' ';
      out << "0\n";
   }
};

inline fun_t expand( fun_t f , unsigned s ) // expand to 4 
{
   extern unsigned short belookup[4][256][4]; // rewrite/be.cpp
   assert( s > 1 );
   if( s == 4 )return f;
   if( s == 3 )return belookup[3][f][3];
   return belookup[3][ belookup[2][f][2] ][3] ;
}

///////////////////////////////////////////////////////////////////////////////
//                          COVERING                                         //
///////////////////////////////////////////////////////////////////////////////

CktCover::CktCover():
   _cover(NULL),_satph(NULL),_nvar(0),_lits(NULL)
{}
CktCover::~CktCover()
{
   if( _cover )
   {
      for( unsigned i = 0 ; i < sz ; ++ i ) 
         delete _cover[i];
      delete [] _cover;
   }
   delete [] _satph;
   delete [] _lits;
}
void CktCover::init( bool phase )
{
   InfoCut::init();

   _cover= new Cover * [sz];
   _lits = new lit_t[sz];
   _satph = new uint8_t[sz];

   memset( _cover , 0 , sizeof( Cover * ) * sz );
   for( unsigned i = 0 ; i < sz; ++i ) _lits[i] = lit_Undef;

   //--- phase
   if( phase )
   {
      memset( _satph , 0 , sizeof( uint8_t ) * sz );
      for( unsigned i = 0 ; i < _po.size() ; ++i )
         _satph[ var(_po[i]) ] |= 1 << sign(_po[i]); 
      for( unsigned i = sz-1; i > _andFlag ; -- i )
      {
         AigNode n = data[i];
         _satph[ var( n.fanin[0] ) ] |= sign( n.fanin[0] ) ? swap01( _satph[i] ) : _satph[i] ;
         _satph[ var( n.fanin[1] ) ] |= sign( n.fanin[1] ) ? swap01( _satph[i] ) : _satph[i] ;
      }
   }else
   {
      memset( _satph , 3 , sizeof( uint8_t ) * sz );
   }
}
void CktCover::covering()
{
   dp_covering();
   optimize();
}
unsigned CktCover::nClause()const
{
   unsigned n = 1;
   unsigned * focnt = _focnt.data;

   for( unsigned i = _andFlag ; i < sz ; ++ i )
   {
      if( focnt[i] == 0 )continue;
      assert( _cover[i] != NULL );
      n+=_cover[i]->ncls();
   }

   return n;
}

///////////////////////////////////////////////////////////////////////////////
//                          CNF ASSIGNMENT                                   //
///////////////////////////////////////////////////////////////////////////////
inline uint8_t cst( uint8_t s ) { return 3-s ;}
void CktCover::assignLit( int var )
{
   _nvar = var;
   l0 = _lits[0] = lit_t( _nvar++,0);
//   for( int i = 1 ; i < sz ; ++ i )
   for( int i = sz-1 ; i > 0 ; -- i )
      if( _lits[i] == lit_Undef && _focnt[i]!=0 )
         _lits[i] = lit_t( _nvar++ , 0 );
}
void CktCover::write_cnf( const char * fname )
{
   SolverFile S( fname , _nvar , nClause() + _po.size() );
   assignCNFPO( S );
   assignCNF( S );
}
void CktCover::assignCNFPO( SATSolver & s )
{
   for( unsigned i = 0 ; i < _po.size() ; ++ i )
      s.addClause( sign(_po[i]) ? ~_lits[ var(_po[i]) ] : _lits[ var(_po[i])] );
}
void CktCover::assignCNF( SATSolver & s )
{  
   unsigned * focnt = _focnt.data;
   for( unsigned i = _andFlag ; i < sz ; ++ i )
   {
      if( focnt[i] == 0 )continue;

      assert( _cover[i] != NULL );
      _cover[i]->dumpClause( i , _lits , s , _satph[i] );
   }

   s.addClause( ~_lits[0] );
}
///////////////////////////////////////////////////////////////////////////////
//                       DYNAMIC PROGRAMMING                                 // 
///////////////////////////////////////////////////////////////////////////////
class ClsDB_c
{
   friend class CktCover;
   struct ClsNPN 
   {
      unsigned short area;
      unsigned short ncls;
      unsigned char * cls;
   };

   NPNPtr       _mixers[1<<(1<<4)];
   vec<ClsNPN>  _clsnpn;
   lit_t         _bucket[4];
   void read( const char * );

   public:
   ClsDB_c( const char * file );
   ~ClsDB_c();

   inline unsigned area( fun_t f ){ return _clsnpn[ _mixers[f].idx ].area; }
   inline unsigned ncls( fun_t f ){ return _clsnpn[ _mixers[f].idx ].ncls; }

   void put_bases( fun_t f , const Cut & c , lit_t * l );
   void dumpClause( fun_t f , lit_t x , SATSolver & S , uint8_t ph);
}ClsDb(CURDIR"dat/cls.db");

///////////////////////////////////////////////////////////////////////////////
//                          MAPPING                                          //
///////////////////////////////////////////////////////////////////////////////

int REFINE_CNT =2;
static float estimate_cost( StrashAIG & ckt , BestCut * best )
{
   float total_cost = 0;
   for( unsigned i = 0 ; i < ckt._po.size() ; ++i ) 
      total_cost += best[ var( ckt._po[i] ) ].cost;
   return total_cost;
}
static void cut_compute( InfoCut & cut )
{
   cut.create_base();
   for( unsigned i = cut._andFlag ; i < cut.size() ; ++i )
      cut.update( i );
}
void CktCover::dp_covering( BestCut * best )
{
   for( int i = sz-1; i >= (int)_andFlag ; -- i )
   {
      BestCut b = best[i];
      Cut c = _cut[i][b.size-1][b.idx];
      _cover[i] = new CutCover( b.size , c );

      if( _focnt[i] != 0 )
      {
         for( int j = 0 ; j < (int)b.size ; ++ j )_focnt[c[j]]+=1;
         rm_cone( i );
      }
   }
}
void CktCover::dp_compute( unsigned short * refcnt , BestCut * best )
{
   for( unsigned node = _andFlag ; node < sz ; ++ node )
   {
      if( _focnt[node] == 0 ) continue;
      float best_cost = 1E38;
      for( int cs = 3 ; cs > 0 ; -- cs )
      {
         CutSet & c = _cut[node][cs];

         for( int i = 0 ; i < c.size ; ++i )
         {
            float cost = ClsDb.area( expand( c[i].fun , cs+1 ) );

            for( int j = 0 ; j <= cs ; ++j )
               cost += best[ c[i][j] ].cost;

            if( cost < best_cost )
            {
               best[node] = BestCut( cs+1 , i , cost); 
               best_cost = cost ;
            }
         }
      }
      assert( best[node].cost < 1E37 );
      assert( best[node].size != 0 );
      (best[node].cost /= (refcnt[node])); // TODO ! useful compemsation
   }
}
void CktCover::dp_refine(BestCut * best ,unsigned short * refcnt )
{
   vec<unsigned> v ;
   for( unsigned i = 0 ; i < _po.size() ; ++i )
   {
      unsigned n = var( _po[i] );
      if( refcnt[ n ] == 0 || n < _andFlag) continue;
      refcnt[n] = 0;
      v.push( n );
   }
   unsigned ptr = 0 ;
   while( ptr != v.size() )
   {
      unsigned idx = v[ ptr ++ ];
      BestCut b = best[idx];
      Cut c = _cut[idx][ b.size -1 ][ b.idx ];
      for( unsigned i = 0 ; i < b.size ; ++ i )
      {
         unsigned n = c[i];
         if( refcnt[n] == 0 || n < _andFlag )continue;
         refcnt[n] = 0;
         v.push( n );
      }
   } 
   for( unsigned i = 0 ; i < v.size() ; ++i )
   {
      unsigned idx = v[i];
      BestCut b = best[idx];
      Cut c = _cut[idx][b.size-1][b.idx];
      for( unsigned j = 0 ; j < b.size ; ++ j )
      {
         unsigned n = c[j];
         if( n < _andFlag )continue;
         refcnt[ n ]+=1;
      }
   }
   for( unsigned i = 0 ; i < _po.size() ; ++i )
   {
      unsigned n = var( _po[i] );
      refcnt[n] +=1 ;
   }
}

void CktCover::dp_covering()
{
   unsigned short * refcnt = new unsigned short [sz];
   BestCut        * best   = new BestCut[sz];
   cut_compute( *this );

   for( unsigned i = 0 ; i < sz ; ++ i ) refcnt[i] = _focnt[i];
   
   dp_compute( refcnt , best);
   for( int i = 0 ; i < REFINE_CNT ; ++ i )
   {
      dp_refine( best , refcnt );
      dp_compute( refcnt , best);
   }
   //cout << "COST = " << (int) estimate_cost( *this , best ) << endl;

   dp_covering( best );

   delete [] best;
   delete [] refcnt;
}
void CktCover::deactivate( unsigned n )
{
   Cover * c = _cover[n];

   for( unsigned i = 0 ; i < c->size() ; ++i )
      if( -- _focnt[ c->at(i) ] == 0 )
          deactivate( c->at(i) );
}
unsigned CktCover::deactivate( unsigned n , unsigned lim , vec<unsigned> & v)
{
   if( lim == 0 )
   {
      v.push( n );
      return 0 ;
   }

   Cover * c = _cover[n];
   unsigned cost = c->area();

   for( unsigned i = 0 ; i < c->size() ; ++i )
      if( -- _focnt[ c->at(i) ] == 0 )
         cost += deactivate( c->at(i) , lim -1 , v);

   return cost;
}
int CktCover::activate( Cover * c , int budget )
{
   budget -= c->area() ;

   int i;
   for( i = 0 ; i < (int)c->size() ; ++ i )
   {
      if( budget <= 0 )break;
      unsigned n = c->at(i);
      if( _focnt[n] == 0 )
         budget = activate( _cover[n] , budget) ;

      _focnt[n] += 1;
   }
   if( budget <= 0 && i!=0 )
   {
      unsigned n = c->at(--i);
      _focnt[n]-=1;

      for( --i ; i >= 0 ; -- i )
      {
         unsigned n = c->at(i);
         _focnt[n] -=1;
         if( _focnt[n] == 0 )
            deactivate( n );
      }
   }

   return budget;
}
void CktCover::deactivate( vec<unsigned> & v )
{
   for( unsigned i = 0 ; i < v.size() ; ++i )
      if( _focnt[ v[i] ] == 0 )
         deactivate( v[i] );
}
void CktCover::optimize()
{
   //assert( check() );
   for( unsigned i = _andFlag; i < sz ; ++ i)
   {
      if( _focnt[ i ] == 0 ) continue;
      for( int s = 1 ; s < 3 ; ++ s ) 
      {
         int lim = _cut[i][s].size;
         for( int j = 0 ; j < lim ; ++ j )
         {
            CutCover * c = new CutCover( s+1 , _cut[i][s][j] );
            if( change( i , c ) == false)
               delete c;
         }
      }
   }
}
void CktCover::bigAnd_covering()
{
   for( unsigned i = sz-1 ; i >= _andFlag ; -- i )
   {
      if( _focnt[i] == 0 ) continue ;
      bigAnd_covering( i );
   }
}
// sign 10  , nosign 01
bool CktCover::isCut_an_and( unsigned k , vec<AigLit> & fanin , uint8_t & sign )
{
   static int db_cut_size[222] = {0,4,3,0,0,2,0};

   CutCover * c = dynamic_cast<CutCover * >(_cover[k]);
   NPNPtr p = ClsDb._mixers[ c->fun ];
   int psz  = db_cut_size[p.idx];
   if( (sign & ( 1<<p.sign )) == 0 )return false;
   sign = ( 1<<p.sign );

   if( psz )
   {
      for( int i = 0 ; i < psz ; ++ i ) 
         fanin.push( toAigLit( c->at( p.in[i].permute ) , p.in[i].sign ));
      return true;
   }
   else return false;
}
bool CktCover::bigAnd_covering( unsigned n )
{
   vec<AigLit> fanin;
   uint8_t sign = 3;
   if( isCut_an_and( n , fanin , sign ) )
   {
      if( fanin.size() != 4 )
      {
         for( unsigned i = 0 ; i < fanin.size() ; ++ i )
            _focnt[ var(fanin[i]) ] +=1;
         return false;
      }
      unsigned ptr = 0 ;
      for( ;ptr!=fanin.size() ; ptr+=1)
      {
         AigLit x = fanin[ptr];
         _focnt[ var(x) ]-=1;
         if( _focnt[var(x)] ) continue ;
         uint8_t tmpsign = 1<<::sign(x);
         if( isCut_an_and( var(x) , fanin , tmpsign ) )
            fanin[ptr] = 0xFFFFFFFF;
      }
      ptr=0;
      for( unsigned i = 0 ; i < fanin.size() ; ++ i )
         if( fanin[i] != 0xFFFFFFFF )
            fanin[ptr++] = fanin[i];
      fanin.downTo( ptr );
      for( unsigned i = 0 ; i < fanin.size() ; ++i )
         _focnt[ var(fanin[i]) ] +=1;

      delete _cover[n];
      _cover[n] = new ANDCover( ptr , fanin.data , sign>>1 );
   }
   return true;
}
bool CktCover::change( unsigned n , Cover * c )
{
   vec<unsigned> v;

   unsigned gain = deactivate( n , 2 , v );
   for( unsigned i = 0 ; i < v.size() ; ++ i ) _focnt[ v[i] ] +=1 ;

   int leave = activate( c , gain);
   if( leave >0 ) 
   {
      delete _cover[n];
      _cover[n] = c;

      for( unsigned i = 0 ; i < v.size() ; ++i ) _focnt[ v[i] ] -=1;
      deactivate( v );
      return true;
   }
   else 
   {
      //activate( _cover[n] , INT_MAX );
      activate( _cover[n] , 0x7FFFFFFF );
      for( unsigned i = 0 ; i < v.size() ; ++i ) _focnt[ v[i] ] -=1;
      return false;
   }
}
///////////////////////////////////////////////////////////////////////////////
//                      CutCover                                             //
///////////////////////////////////////////////////////////////////////////////
CutCover::CutCover( unsigned short s ,const Cut & c):Cut(c),sz(s)
{
   fun = expand( fun , s );
   for( ; s != 4 ; ++s ) node[s] = 0; // avoid put uninitial value at base 
}
unsigned CutCover::area(){return ClsDb.area( fun );}
unsigned CutCover::ncls(){return ClsDb.ncls( fun );}
void CutCover::dumpClause( unsigned n , lit_t * l , SATSolver & S , uint8_t ph )
{
   ClsDb.put_bases ( fun , *this , l );
   ClsDb.dumpClause( fun , l[n] , S , ph );
}
void ANDCover::dumpClause( unsigned n , lit_t * l , SATSolver & S , uint8_t ph )
{
   if( ph == 0 )return;

   lit_t target = sign?~l[n]:l[n];
   vec<lit_t> v;    v.push( target );
   target = ~target;
   for( unsigned i = 0 ; i < sz ; ++ i )
   {
      lit_t x  = ::sign(dat[i]) ? ~l[var(dat[i])] : l[var(dat[i])] ;
      v.push( ~x );
      S.addClause( x , target );
   }
   S.addClause( v );
}
///////////////////////////////////////////////////////////////////////////////
//                      CutCover                                             //
///////////////////////////////////////////////////////////////////////////////
ClsDB_c::ClsDB_c( const char * file ){read( file );}
void ClsDB_c::read( const char * file )
{
   std::ifstream db( file , std::ios::binary | std::ios::in );
   db.read( (char*)_mixers , sizeof( NPNPtr)* (1<<(1<<4)) );
   _clsnpn.resize( 222 );
   for( int i = 0 ; i < 222 ; ++i )
   {    
      ClsNPN & c = _clsnpn[i];
      db.read( (char *)& c.area , sizeof(unsigned short ) );
      db.read( (char *)& c.ncls , sizeof(unsigned short ) );
      c.cls = new unsigned char [c.ncls+2];
      db.read( (char *) c.cls , c.ncls+2 );
   }
}
ClsDB_c::~ClsDB_c()
{
   for(int i = 0 ; i < 222 ; ++i )
      delete [] _clsnpn[i].cls ;
}
void ClsDB_c::put_bases( fun_t f , const Cut & c , lit_t * l )
{
   NPNPtr mixer     = _mixers[f];
   for( int i = 0 ; i < 4 ; ++ i )
   {
      lit_t y = l[ c[ mixer.in[i].permute] ];
      if( c[ mixer.in[i].permute] == 0 )continue; 
      _bucket[ i ] = lit_t( y.var() , mixer.in[i].sign ^ y.sign() );
   }
}
void ClsDB_c::dumpClause( fun_t f , lit_t x , SATSolver & S , uint8_t ph )
{
   if( ph == 0 )return;

   NPNPtr mixer = _mixers[f];
   if( mixer.sign == 0 ) x = ~x;
   else ph = swap01( ph );
   vec<lit_t> v(4);

   const unsigned char * p = _clsnpn[ mixer.idx ].cls;
   for( int k = 0 ; k < 2 ; ++k )
   {
      while( (*p) != 0xFF )
      {
         v.clear();

         char c = *p++;
         for( int i = 0 ; i < 4 ; ++i , c=c>>2)
         {
            switch( c & 3 )
            {
               case 3 : assert(false);break;
               case 2 : break;
               case 0 : v.push(_bucket[i]);  assert(_bucket[i]!=lit_Undef);break;
               case 1 : v.push(~_bucket[i]); assert(_bucket[i]!=lit_Undef);break;
            }
         }
         v.push( x );
         if( ph & 1 || v.size() <= 2) 
            if( slim(v) )
               S.addClause( v );
      }
      p++;
      x=~x;
      ph = swap01( ph );
   }
}
void CktCover::const_propagation( )
{
   uint8_t * csts = new uint8_t [sz];
   memset( csts , 0 , sizeof( uint8_t ) * sz );

   for( unsigned i = 0 ; i < _po.size() ; ++ i )
      csts[ var(_po[i]) ] = cst( sign( _po[i] ) );
   for( unsigned i = sz-1; i >= _andFlag ; --i )
   {
      if( csts[i] == 3 )
      {
         AigNode n = data[i];
         csts[ var(n.fanin[0]) ] = cst( sign( n.fanin[0] ) );
         csts[ var(n.fanin[1]) ] = cst( sign( n.fanin[1] ) );
      }
   }
   for( unsigned i = _andFlag ; i < sz ; ++ i )
   {
      if( csts[i]&2 ) continue;
      AigNode n = data[i];
      uint8_t v[2] = { csts[ var(n.fanin[0]) ] ^ sign( n.fanin[0] ) , csts[ var(n.fanin[1]) ] ^ sign( n.fanin[1] ) };
      if( v[0] == 2 || v[1] == 2 ) csts[i] = 2 ;
      if( v[0] == 3 && v[1] == 3 ) csts[i] = 3;
   }
   unsigned cnt = 0;
   for( unsigned i = 1 ; i < sz ; ++ i )
      if( csts[i] & 2 )
      {
         cnt += 1 ;
         _lits[ i ] = ((csts[i]==2)?l0:~l0);
      }

   // write cnf will report value 1  

   delete [] csts;
}
Lit CktCover::getLit( AigLit x )const
{
   return _lits[ var(x) ] ^ sign(x) ;
}

#if 0
void CktCover::report_covering()const
{
   unsigned covered=0;
   unsigned area = 0;
   unsigned ncls=0;

   for( unsigned i = _andFlag ; i < sz ; ++ i) 
   {
      if( _focnt[i] == 0 )continue;

      if( _cover[i] != NULL )
      {
         covered+=1;
         area += _cover[i]->area();
         ncls += _cover[i]->ncls();
      }else 
         area += 7,
         ncls += 3;
   }

   cout << "c cnfcnt : "<<sz-_andFlag << "\t" << covered<<"\t"<< area <<"\t"<< ncls <<endl;
}
bool CktCover::check()
{
   vec<unsigned> v;

   for( int sel = 0 ; sel < 2 ; ++ sel )
   {
      vec<unsigned> & po = sel?_ppo:_po;
      for( unsigned i = 0 ; i < po.size() ; ++i )
         if( --_focnt[ var(po[i]) ] == 0 )
            v.push( var( po[i] ) );
   }
   unsigned ptr = 0;
   while( ptr != v.size() )
   {
      unsigned n = v[ptr++];
      if( _focnt[n] != 0 ){cout<<"Error focnt " << n << " = " << _focnt[n] << endl; return false;}
      Cover * c = _cover[n];
      for( unsigned j = 0 ; j < c->size() ; ++ j )
         if( -- _focnt[ c->at(j) ] == 0 )
            v.push( c->at(j) );
   }
   for( unsigned i = 0 ; i < _andFlag;  ++ i ) 
   {
      if( _focnt[i] != 1 ) {
         cout << "Error focnt "<< i << " = "<< _focnt[i] <<endl;
         return false;
      }
   }
   for( unsigned i = 0 ; i < v.size() ; ++ i )
   {
      Cover * c = _cover[ v[i] ];
      for( unsigned j = 0 ; j < c->size() ; ++ j )
         _focnt[ c->at(j) ] +=1; 
   }
   for( int sel = 0 ; sel < 2 ; ++ sel )
   {
      vec<unsigned> & po = sel?_ppo:_po;
      for( unsigned i = 0 ; i < po.size() ; ++i )
         _focnt[ var(po[i]) ] +=1;
   }
   return true;
}

#endif
